#!/usr/bin/env python
# -*- coding: utf-8 -*-
# https://cloud.tencent.com/developer/article/1111575
# gensim modules
import random
import numpy
import numpy as np
import gensim
# from gensim import utils
import torch
from torch import tensor
from gensim.models.doc2vec import TaggedDocument
# from gensim.models import Doc2Vec
# from sklearn.linear_model import LogisticRegression# classifier
# from sklearn.model_selection import train_test_split
import logging as log
# import sys
# from sklearn.metrics import roc_curve, auc
# import matplotlib.pyplot as plt

#


"""list2docvec.py
#功能：list的向量化表示
输入形式：list=[str,];注意不全为下划线连接，否则构造词汇表错误
输入：[' derivationally related form', ' derivationally related form']
输出向量（tensor）：两个模型上的向量表示train_arrays_dm:(4, size),train_arrays_dbow:(4, size)（vector_size由size确定）
输出日志：./data/list2vec.log
"""
#参数设置

data_log="./data/list2vec.log"
#RuntimeError: you must first build vocabulary before training the model
# 创建TaggedDocument对象-类似字典形式{标签：str}
def labelizeReviews(reviews, label_type):
    #对列表逐一生成标签
    #生成类似字典：{label_type_i:reviews}，reviews-list,label_type-str
    for i, v in enumerate(reviews):
        label = '%s_%s' % (label_type, i)#标签形式
        yield TaggedDocument(gensim.utils.simple_preprocess(v, max_len=100), [label])

#打乱词汇表
def sentences_perm(sentences):
    shuffled = list(sentences)
    random.shuffle(shuffled)
    return (shuffled)

def list2dv(inlist,size=200,epochs=10):
    """
    功能：list文本的doc2vec向量化表示；
    形式：[str,]->torch.Size(len(list).size)
    :param inlist: 形式=[str,]

    :param size: size = 200#对应gensim.models.Doc2Vec中的vector_size（每行文本对应的向量维度）
    :return:dm和dbow两个模型处理后的向量-device:cuda:0,type:<class 'torch.Tensor'>,arrays_dm:torch.Size([4, 200]),arrays_dbow:torch.Size([4, 200])
    调用：
    inlist=['_derivationally_related_form', '_derivationally_related_form', '_hypernym', '_member_meronym']
    spalist=list_replace(inlist,'_',' ')#替代
    train_arrays_dm,train_arrays_dbow=list2dv(spalist)
    print("device:{},type:{},arrays_dm:{},arrays_dbow:{}".format(train_arrays_dm.device,type(train_arrays_dm),train_arrays_dm.shape,train_arrays_dbow.shape))#
    """


    x_train_tag = list(labelizeReviews(inlist, 'train'))

    # 下面我们实例化两个 Doc2Vec 模型，DM 和 DBOW
    # 实例化 DM 和 DBOW 模型
    model_dm = gensim.models.Doc2Vec(min_count=1, window=10, vector_size=size, sample=1e-3, negative=5, workers=3,
                                     epochs=10)
    model_dbow = gensim.models.Doc2Vec(min_count=1, window=10, vector_size=size, sample=1e-3, negative=5, dm=0, workers=3,
                                       epochs=10)
    # 对所有评论创建词汇表（正、负、未标记），词汇表可看作字典形式
    alldata = x_train_tag
    model_dm.build_vocab(alldata)#创建词汇表-模型1
    model_dbow.build_vocab(alldata)#创建词汇表-模型2

    for epoch in range(epochs):#训练2个模型
        log.info('EPOCH: {}'.format(epoch))
        model_dm.train(sentences_perm(alldata), total_examples=model_dm.corpus_count, epochs=1)
        model_dbow.train(sentences_perm(alldata), total_examples=model_dbow.corpus_count, epochs=1)

    # 取向量：用标签或者内容
    # 取向量（第一种方法）:用标签
    #取得训练集向量-连接2个模型
    train_arrays_dm = numpy.zeros((len(inlist), size))
    train_arrays_dbow = numpy.zeros((len(inlist), size))
    for i in range(len(inlist)):
        tag = 'train_' + str(i)#标签
        train_arrays_dm[i] = model_dm.dv[tag]#由标签取训练后的向量-DB 模型
        train_arrays_dbow[i] = model_dbow.dv[tag]#由标签取训练后的向量-DBOW 模型
    train_arrays = np.hstack((train_arrays_dm, train_arrays_dbow))#连接2个模型的向量#print("output array's shape:",train_arrays.shape)#
    #ndarray处理为tensor
    train_arrays_dm=torch.from_numpy(train_arrays_dm).type(torch.FloatTensor).cuda() # array to tensor
    train_arrays_dbow = torch.from_numpy(train_arrays_dbow).type(torch.FloatTensor).cuda()   # array to tensor
    train_arrays = torch.from_numpy(train_arrays).type(torch.FloatTensor).cuda()   # array to tensor
    # print("device:{},type:{},arrays_dm:{},arrays_dbow:{}".format(train_arrays_dm.device,type(train_arrays_dm),train_arrays_dm.shape,train_arrays_dbow.shape))#
    return train_arrays_dm, train_arrays_dbow

def list2dv_dict(inlist,size=20,epochs=10):
    """
    功能：list文本的doc2vec字典向量化表示；
    形式：list->{inlist[i]:tensor(size)}
    参数：
        :param inlist: 形式=[str,]

        :param size: size = 200#对应gensim.models.Doc2Vec中的vector_size（每行文本对应的向量维度）
        :return:dm和dbow两个模型处理后的向量-device:cuda:0,type:<class 'torch.Tensor'>,arrays_dm:torch.Size([4, 200]),arrays_dbow:torch.Size([4, 200])
    调用：
        inlist=['_derivationally_related_form', '_derivationally_related_form', '_hypernym', '_member_meronym']
        spalist=list_replace(inlist,'_',' ')#替代
        train_arrays_dm,train_arrays_dbow=list2dv_dict(spalist)
        print(f'dm:{list(train_arrays_dm.values())[0].shape}')#torch.Size([size])
    """
    x_train_tag = list(labelizeReviews(inlist, 'train'))

    # 下面我们实例化两个 Doc2Vec 模型，DM 和 DBOW
    # 实例化 DM 和 DBOW 模型
    model_dm = gensim.models.Doc2Vec(min_count=1, window=10, vector_size=size, sample=1e-3, negative=5, workers=3,
                                     epochs=10)
    model_dbow = gensim.models.Doc2Vec(min_count=1, window=10, vector_size=size, sample=1e-3, negative=5, dm=0, workers=3,
                                       epochs=10)
    # 对所有评论创建词汇表（正、负、未标记），词汇表可看作字典形式
    alldata = x_train_tag
    model_dm.build_vocab(alldata)#创建词汇表-模型1
    model_dbow.build_vocab(alldata)#创建词汇表-模型2

    for epoch in range(epochs):#训练2个模型
        log.info('EPOCH: {}'.format(epoch))
        model_dm.train(sentences_perm(alldata), total_examples=model_dm.corpus_count, epochs=1)
        model_dbow.train(sentences_perm(alldata), total_examples=model_dbow.corpus_count, epochs=1)

    # 取向量：用标签或者内容
    # 取向量（第一种方法）:用标签
    #取得训练集向量-连接2个模型
    train_dict_dm = {}
    train_dict_dbow = {}
    for i in range(len(inlist)):
        tag = 'train_' + str(i)#标签
        train_dict_dm[inlist[i]] = torch.tensor(model_dm.dv[tag])#由标签取训练后的向量-DB 模型
        train_dict_dbow[inlist[i]] = torch.tensor(model_dbow.dv[tag])#由标签取训练后的向量-DBOW 模型

    return train_dict_dm, train_dict_dbow


def list_replace(inlist,old,new):
    """
    替换列表中的字符
    :param inlist: list
    :param old: str,旧字符
    :param new: str,新字符
    :return: list
    调用：
    inlist=['_derivationally_related_form', '_derivationally_related_form']
    nlist=list_replace(inlist,'_',' ')
    """
    nlist = []
    for i in inlist:
        nlist.append(i.replace('_',' '))
    # print(nlist)
    return nlist


